import os
import pathlib
import sys
import json
import matplotlib.pyplot as plt
import torch
import inspect
from collections import OrderedDict
from typing import Any, Callable, Dict
import zipfile

__all__ = ["PYTHONPATH", "zip_utils"]


def have_ipython():
    return "get_ipython" in globals()


class Inspector(object):
    def __init__(self, base_class=None):
        self.base_name = base_class.__class__.__name__+"." if base_class else ""
        self.params: Dict[str, Dict[str, Any]] = {}

    def inspect(self, func: Callable):
        params = inspect.signature(func).parameters
        params = OrderedDict(params)
        self.params[func.__name__] = params

    def distribute(self, func_name, kwargs: Dict[str, Any]):
        out = {}
        for key, param in self.params[func_name].items():
            data = kwargs.get(key, inspect.Parameter.empty)
            if data is inspect.Parameter.empty:
                if param.default is inspect.Parameter.empty:
                    raise TypeError(
                        f'{self.base_name}{func_name} Required parameter {key} is empty.')
                data = param.default
            out[key] = data
        return out


def write_dict_to_file(d: dict, path: str):
    with open(path, encoding="utf8", mode="w") as f:
        json.dump(d, f)


def show(images, labels, rows, cols, colors=None, figsize=None):
    fig, ax = plt.subplots(rows, cols, figsize=figsize)
    ax = ax.flatten()
    for i in range(rows * cols):
        ax[i].imshow(images[i])
        ax[i].set_title(labels[i], c=colors[i] if colors else 'g')
        ax[i].axis("off")
    plt.tight_layout()
    plt.show()


def show_batch(dataloader, rows=4, cols=4, figsize=None):
    images, labels = [], []
    nums = rows * cols
    idx_to_class = {v: k for k, v in dataloader.dataset.class_to_idx.items()}
    for img, label in dataloader:
        if len(images) >= nums:
            break
        images.extend(torch.einsum("b c h w -> b h w c", img).cpu().unbind())
        labels.extend([idx_to_class[l] for l in label.tolist()])
    show(images, labels, rows, cols, figsize=figsize)


def predict_batch(dataloader, model, device, dtype, rows=4, cols=4, figsize=None):
    images, labels, predicts = [], [], []
    nums = rows * cols
    idx_to_class = {v: k for k, v in dataloader.dataset.class_to_idx.items()}
    model.to(device, dtype)
    model.eval()
    for img, label in dataloader:
        if len(images) >= nums:
            break
        images.extend(torch.einsum("b c h w -> b h w c", img).cpu().unbind())
        labels.extend([idx_to_class[l] for l in label.tolist()])
        o = model(img.to(device, dtype)).argmax(dim=-1).tolist()
        predicts.extend([idx_to_class[l] for l in o])
    out = [f"类别:{l}\n预测:{p}" for l, p in zip(labels, predicts)]
    color = ["g" if l == p else "r" for l, p in zip(labels, predicts)]
    show(images, out, rows, cols, color, figsize=figsize)


def dict_as_table(config):
    def string(msg): return msg.__name__ if hasattr(
        msg, "__name__") else str(msg)
    config = {str(k): string(v) for k, v in config.items()}
    max_keys_len = max(map(len, config.keys()))
    max_values_len = max(map(len, config.values()))
    n = max_keys_len + max_values_len + 6 + 2
    out_str = f"{'-' * n}\n"
    for k, v in config.items():
        out_str += f"|  {k:<{max_keys_len}}  {v:>{max_values_len}}  |\n"
        out_str += f"{'-' * n}\n"
    return out_str


class zip_utils:
    @staticmethod
    def zip_dir(dir, zip_name):
        pathlib.Path(zip_name).parent.mkdir(parents=True, exist_ok=True)
        zip_name = zip_name if zip_name.endswith('.zip') else zip_name + '.zip'
        with zipfile.ZipFile(zip_name, 'w') as z:
            for root, dirs, files in os.walk(dir):
                for file in files:
                    z.write(os.path.join(root, file))
        return zip_name

    @staticmethod
    def create_zip_from_list(zip_name, file_name_list):
        zip_name = zip_name if zip_name.endswith('.zip') else zip_name + '.zip'
        with zipfile.ZipFile(zip_name, 'w') as z:
            for file_name in file_name_list:
                z.write(file_name)

    @staticmethod
    def create_zip_from_pattern(zip_name, pattern):
        import glob
        file_name_list = glob.glob(pattern)
        return zip_utils.create_zip_from_list(zip_name, file_name_list)


class PYTHONPATH_ADDER:
    def __lshift__(self, other):
        if other not in sys.path:
            sys.path.append(other)
        return self

    def __contains__(self, item):
        return item in sys.path


PYTHONPATH = PYTHONPATH_ADDER()
